#!/bin/bash

# IO定时监控脚本
# 用于在后台定期检查IO状态并在发现异常时发出警告

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 默认参数
INTERVAL=300  # 默认检查间隔为5分钟
THRESHOLD=90  # 默认IO使用率警告阈值为90%
LOG_FILE="/var/log/io-monitor.log"
NOTIFY=1      # 默认启用桌面通知
SERVICE_MODE=0 # 默认不以服务模式运行

# 帮助信息
show_help() {
    echo "IO定时监控脚本 - 在后台定期检查IO状态并在发现异常时发出警告"
    echo
    echo "用法: $0 [选项]"
    echo
    echo "选项:"
    echo "  -i, --interval <秒>    检查间隔，默认300秒(5分钟)"
    echo "  -t, --threshold <百分比> IO使用率警告阈值，默认90%"
    echo "  -l, --log <文件>       日志文件路径，默认/var/log/io-monitor.log"
    echo "  -n, --no-notify        禁用桌面通知"
    echo "  -s, --service          以服务模式运行（后台运行）"
    echo "  -h, --help             显示此帮助信息"
    echo
    echo "示例:"
    echo "  $0 -i 60 -t 80         每60秒检查一次，IO使用率超过80%时发出警告"
    echo "  $0 -s -l /tmp/io.log   以服务模式运行，日志保存到/tmp/io.log"
}

# 解析参数
while [[ $# -gt 0 ]]; do
    case $1 in
        -i|--interval)
            INTERVAL="$2"
            shift 2
            ;;
        -t|--threshold)
            THRESHOLD="$2"
            shift 2
            ;;
        -l|--log)
            LOG_FILE="$2"
            shift 2
            ;;
        -n|--no-notify)
            NOTIFY=0
            shift
            ;;
        -s|--service)
            SERVICE_MODE=1
            shift
            ;;
        -h|--help)
            show_help
            exit 0
            ;;
        *)
            echo -e "${RED}未知选项: $1${NC}"
            show_help
            exit 1
            ;;
    esac
done

# 检查依赖
check_dependencies() {
    local missing_deps=()
    
    if ! command -v iostat &> /dev/null; then
        missing_deps+=("sysstat")
    fi
    
    if [ "$NOTIFY" -eq 1 ] && ! command -v notify-send &> /dev/null; then
        missing_deps+=("libnotify-bin")
    fi
    
    if [ ${#missing_deps[@]} -gt 0 ]; then
        echo -e "${RED}缺少依赖:${NC} ${missing_deps[*]}"
        echo "请安装缺少的依赖: sudo apt install ${missing_deps[*]}"
        exit 1
    fi
}

# 日志函数
log_message() {
    local level="$1"
    local message="$2"
    local timestamp=$(date "+%Y-%m-%d %H:%M:%S")
    
    echo "[$timestamp] [$level] $message" >> "$LOG_FILE"
    
    # 如果不是服务模式，也输出到控制台
    if [ "$SERVICE_MODE" -eq 0 ]; then
        case "$level" in
            INFO)
                echo -e "${BLUE}[$timestamp] [$level] $message${NC}"
                ;;
            WARNING)
                echo -e "${YELLOW}[$timestamp] [$level] $message${NC}"
                ;;
            ERROR|CRITICAL)
                echo -e "${RED}[$timestamp] [$level] $message${NC}"
                ;;
            SUCCESS)
                echo -e "${GREEN}[$timestamp] [$level] $message${NC}"
                ;;
            *)
                echo "[$timestamp] [$level] $message"
                ;;
        esac
    fi
}

# 发送桌面通知
send_notification() {
    local title="$1"
    local message="$2"
    local urgency="$3"
    
    if [ "$NOTIFY" -eq 1 ] && command -v notify-send &> /dev/null; then
        # 获取当前用户
        local current_user=$(logname 2>/dev/null || echo $SUDO_USER)
        if [ -z "$current_user" ]; then
            current_user=$(who | awk '{print $1}' | head -1)
        fi
        
        if [ -n "$current_user" ]; then
            # 获取当前用户的DBUS会话总线地址
            local dbus_address=$(grep -z DBUS_SESSION_BUS_ADDRESS /proc/$(pgrep -u $current_user gnome-session | head -1)/environ 2>/dev/null)
            
            if [ -n "$dbus_address" ]; then
                # 以当前用户身份发送通知
                sudo -u $current_user DBUS_SESSION_BUS_ADDRESS="${dbus_address#*=}" notify-send --urgency="$urgency" "$title" "$message"
            else
                log_message "WARNING" "无法获取DBUS会话地址，无法发送桌面通知"
            fi
        else
            log_message "WARNING" "无法确定当前用户，无法发送桌面通知"
        fi
    fi
}

# 检查IO状态
check_io_status() {
    # 获取IO统计信息
    local io_stats=$(iostat -xd 1 1 | grep -v "^$" | grep -A 1 "^Device" | tail -n +2)
    local high_io_detected=0
    local high_io_devices=""
    
    # 分析IO统计信息
    echo "$io_stats" | while read line; do
        if [ -n "$line" ]; then
            local device=$(echo "$line" | awk '{print $1}')
            local util=$(echo "$line" | awk '{print $NF}')
            
            # 检查是否为数字
            if [[ "$util" =~ ^[0-9]+(\.[0-9]+)?$ ]]; then
                local util_int=${util%.*}
                
                # 检查是否超过阈值
                if [ "$util_int" -ge "$THRESHOLD" ]; then
                    log_message "WARNING" "设备 $device IO使用率过高: ${util}%"
                    high_io_detected=1
                    high_io_devices="$high_io_devices$device (${util}%), "
                fi
            fi
        fi
    done
    
    # 如果检测到高IO使用率，发送通知
    if [ "$high_io_detected" -eq 1 ]; then
        # 移除最后的逗号和空格
        high_io_devices=${high_io_devices%, }
        
        # 获取等待IO的进程
        local io_wait_procs=$(ps aux | awk '$8 ~ /D/ {print $0}' | head -3)
        local proc_info=""
        
        if [ -n "$io_wait_procs" ]; then
            proc_info="\n\n等待IO的进程:\n$io_wait_procs"
        fi
        
        # 发送通知
        send_notification "IO使用率警告" "检测到高IO使用率:\n$high_io_devices$proc_info" "critical"
        
        return 1
    fi
    
    return 0
}

# 主监控循环
monitor_loop() {
    log_message "INFO" "IO监控服务已启动，检查间隔: ${INTERVAL}秒，警告阈值: ${THRESHOLD}%"
    
    while true; do
        # 检查IO状态
        if ! check_io_status; then
            # 如果检测到高IO使用率，增加检查频率
            sleep 60
        else
            # 正常间隔
            sleep "$INTERVAL"
        fi
    done
}

# 创建系统服务
create_service() {
    local service_file="/etc/systemd/system/io-monitor.service"
    local script_path=$(readlink -f "$0")
    
    cat > "$service_file" << EOF
[Unit]
Description=IO Monitor Service
After=network.target

[Service]
Type=simple
ExecStart=$script_path -s -i $INTERVAL -t $THRESHOLD -l $LOG_FILE
Restart=always
RestartSec=10

[Install]
WantedBy=multi-user.target
EOF
    
    systemctl daemon-reload
    systemctl enable io-monitor.service
    systemctl start io-monitor.service
    
    log_message "SUCCESS" "IO监控服务已创建并启动"
    echo -e "${GREEN}IO监控服务已创建并启动${NC}"
    echo "使用以下命令管理服务:"
    echo "  查看状态: systemctl status io-monitor.service"
    echo "  停止服务: systemctl stop io-monitor.service"
    echo "  禁用服务: systemctl disable io-monitor.service"
}

# 主函数
main() {
    # 检查是否以root权限运行
    if [[ $EUID -ne 0 ]]; then
        echo -e "${RED}此脚本需要root权限运行${NC}"
        echo "请使用 sudo $0 重新运行"
        exit 1
    fi
    
    # 检查依赖
    check_dependencies
    
    # 创建日志目录
    mkdir -p $(dirname "$LOG_FILE")
    
    # 询问是否安装为系统服务
    if [ "$SERVICE_MODE" -eq 0 ]; then
        read -p "是否将IO监控安装为系统服务? (y/n): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            create_service
            exit 0
        fi
    fi
    
    # 如果是服务模式，在后台运行
    if [ "$SERVICE_MODE" -eq 1 ]; then
        monitor_loop &
        log_message "INFO" "IO监控服务已在后台启动，进程ID: $!"
        exit 0
    else
        # 前台运行
        echo -e "${BLUE}===== IO定时监控 =====${NC}"
        echo "检查间隔: ${INTERVAL}秒"
        echo "警告阈值: ${THRESHOLD}%"
        echo "日志文件: ${LOG_FILE}"
        echo "桌面通知: $([ "$NOTIFY" -eq 1 ] && echo "已启用" || echo "已禁用")"
        echo
        echo "按 Ctrl+C 停止监控"
        echo
        
        monitor_loop
    fi
}

# 脚本入口
main